"""
测试AI工具的客户端
专门演示和测试集成的Qwen AI功能
"""
import asyncio
import json
from mcp_client import MCPClient


class AIToolsTester:
    """AI工具测试器"""
    
    def __init__(self, server_url: str = "http://localhost:8000"):
        self.client = MCPClient(server_url)
    
    async def run_ai_tests(self):
        """运行AI工具测试"""
        print("🤖 MCP AI工具测试")
        print("="*60)
        
        # 连接服务器
        if not await self.client.connect():
            print("❌ 无法连接到服务器，请确保AI服务器正在运行")
            print("提示: python mcp_ai_server.py")
            return
        
        # 展示AI能力
        await self.show_ai_capabilities()
        
        # 测试各种AI工具
        await self.test_ai_chat()
        await self.test_ai_code_review()
        await self.test_ai_documentation()
        await self.test_ai_summarize()
        await self.test_ai_translate()
        
        # 综合AI工作流演示
        await self.demo_ai_workflow()
        
        await self.client.disconnect()
    
    async def show_ai_capabilities(self):
        """展示AI能力"""
        print("\n📋 AI工具能力清单:")
        print("-" * 40)
        
        if self.client.capabilities:
            ai_tools = [tool for tool in self.client.capabilities.tools if tool.name.startswith('ai_')]
            
            for tool in ai_tools:
                print(f"🔧 {tool.name}: {tool.description}")
                for param in tool.parameters:
                    required = "✓" if param.required else "○"
                    default = f" (默认: {param.default})" if param.default else ""
                    print(f"   {required} {param.name}: {param.description}{default}")
                print()
    
    async def test_ai_chat(self):
        """测试AI对话功能"""
        print("\n" + "="*30)
        print("🗣️ AI对话测试")
        print("="*30)
        
        test_cases = [
            {
                "message": "请介绍一下MCP协议的核心概念和三个基元",
                "description": "基础问答测试"
            },
            {
                "message": "如何设计一个安全的AI系统？",
                "system_prompt": "你是一位AI安全专家，请提供专业的建议",
                "description": "系统提示词测试"
            },
            {
                "message": "写一个Python装饰器的示例",
                "temperature": 0.3,
                "description": "代码生成测试（低温度）"
            }
        ]
        
        for i, test_case in enumerate(test_cases, 1):
            print(f"\n{i}️⃣ {test_case['description']}:")
            print(f"问题: {test_case['message']}")
            
            try:
                result = await self.client.call_tool("ai_chat", test_case)
                if result.success:
                    content = result.result.get("content", "")
                    model = result.result.get("model", "")
                    usage = result.result.get("usage", {})
                    
                    print(f"✅ 模型: {model}")
                    print(f"📊 Token使用: {usage}")
                    print(f"💬 回复: {content[:300]}{'...' if len(content) > 300 else ''}")
                else:
                    print(f"❌ 失败: {result.error}")
            except Exception as e:
                print(f"❌ 异常: {e}")
    
    async def test_ai_code_review(self):
        """测试AI代码审查功能"""
        print("\n" + "="*30)
        print("🔍 AI代码审查测试")
        print("="*30)
        
        test_codes = [
            {
                "code": """
def login(username, password):
    # 不安全的SQL查询
    query = f"SELECT * FROM users WHERE username='{username}' AND password='{password}'"
    result = db.execute(query)
    if result:
        session['user'] = username  # 会话管理不当
        return True
    return False
""",
                "language": "python",
                "focus": "安全性漏洞",
                "description": "SQL注入和会话安全问题"
            },
            {
                "code": """
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)
""",
                "language": "python", 
                "focus": "性能优化",
                "description": "斐波那契递归性能问题"
            }
        ]
        
        for i, test_case in enumerate(test_codes, 1):
            print(f"\n{i}️⃣ {test_case['description']}:")
            
            try:
                result = await self.client.call_tool("ai_code_review", test_case)
                if result.success:
                    content = result.result.get("content", "")
                    print(f"✅ 审查结果: {content[:400]}{'...' if len(content) > 400 else ''}")
                else:
                    print(f"❌ 失败: {result.error}")
            except Exception as e:
                print(f"❌ 异常: {e}")
    
    async def test_ai_documentation(self):
        """测试AI文档生成功能"""
        print("\n" + "="*30)
        print("📚 AI文档生成测试")
        print("="*30)
        
        test_cases = [
            {
                "content": "class MCPServer:\n    def __init__(self, host='localhost', port=8000):\n        self.host = host\n        self.port = port\n    \n    async def start(self):\n        # 启动服务器逻辑\n        pass",
                "doc_type": "API文档",
                "style": "详细",
                "description": "类文档生成"
            },
            {
                "content": "/api/tools/call - POST请求，调用指定工具",
                "doc_type": "接口文档",
                "style": "简洁",
                "description": "API接口文档"
            }
        ]
        
        for i, test_case in enumerate(test_cases, 1):
            print(f"\n{i}️⃣ {test_case['description']}:")
            
            try:
                result = await self.client.call_tool("ai_documentation", test_case)
                if result.success:
                    content = result.result.get("content", "")
                    print(f"✅ 生成文档: {content[:400]}{'...' if len(content) > 400 else ''}")
                else:
                    print(f"❌ 失败: {result.error}")
            except Exception as e:
                print(f"❌ 异常: {e}")
    
    async def test_ai_summarize(self):
        """测试AI内容总结功能"""
        print("\n" + "="*30)
        print("📄 AI内容总结测试")
        print("="*30)
        
        long_text = """
        MCP (Model Context Protocol) 是一个标准化的协议，用于AI模型与外部世界的动态交互。
        本项目实现了一个完整的MCP系统，包括MCP Server、MCP Client以及三个核心基元的完整实现。
        
        三个核心基元包括：
        1. 工具(Tools) - 模型控制：AI可以自主决定调用的功能，如发送通知、创建文件、执行命令等。
        2. 资源(Resources) - 应用控制：提供给AI的结构化数据流，如系统日志、天气信息、数据库内容等。
        3. 提示(Prompts) - 用户控制：可重用的指令模板，如代码审查模板、文档生成模板等。
        
        这种分层控制机制确保了AI系统的安全性和可控性，同时为AI智能体提供了强大的扩展能力。
        通过标准化的接口，AI模型能够安全、可控地与外部世界进行动态交互，为构建未来的AI智能体架构提供了重要的技术基础。
        """
        
        test_cases = [
            {
                "text": long_text,
                "length": "简短",
                "description": "简短总结测试"
            },
            {
                "text": long_text,
                "length": "中等",
                "focus": "技术架构",
                "description": "中等长度，重点总结"
            }
        ]
        
        for i, test_case in enumerate(test_cases, 1):
            print(f"\n{i}️⃣ {test_case['description']}:")
            
            try:
                result = await self.client.call_tool("ai_summarize", test_case)
                if result.success:
                    content = result.result.get("content", "")
                    print(f"✅ 总结结果: {content}")
                else:
                    print(f"❌ 失败: {result.error}")
            except Exception as e:
                print(f"❌ 异常: {e}")
    
    async def test_ai_translate(self):
        """测试AI翻译功能"""
        print("\n" + "="*30)
        print("🌐 AI翻译测试")
        print("="*30)
        
        test_cases = [
            {
                "text": "MCP协议实现了AI智能体与外部世界的标准化交互",
                "target_lang": "English",
                "description": "中译英"
            },
            {
                "text": "The Model Context Protocol enables standardized interaction between AI agents and the external world",
                "target_lang": "中文",
                "description": "英译中"
            }
        ]
        
        for i, test_case in enumerate(test_cases, 1):
            print(f"\n{i}️⃣ {test_case['description']}:")
            print(f"原文: {test_case['text']}")
            
            try:
                result = await self.client.call_tool("ai_translate", test_case)
                if result.success:
                    content = result.result.get("content", "")
                    target_lang = result.result.get("target_language", "")
                    print(f"✅ 翻译到{target_lang}: {content}")
                else:
                    print(f"❌ 失败: {result.error}")
            except Exception as e:
                print(f"❌ 异常: {e}")
    
    async def demo_ai_workflow(self):
        """演示AI工具的综合工作流"""
        print("\n" + "="*40)
        print("🔄 AI工具综合工作流演示")
        print("="*40)
        print("场景: AI代码助手完整工作流")
        print("1. AI分析用户需求")
        print("2. AI审查代码质量")
        print("3. AI生成改进文档") 
        print("4. AI总结工作成果")
        
        # 步骤1: AI分析需求
        print("\n📋 步骤1: AI需求分析")
        need_analysis = await self.client.call_tool("ai_chat", {
            "message": "我需要优化一个用户认证系统，目前有安全性和性能问题，请给出分析建议",
            "system_prompt": "你是一位资深的系统架构师，请提供专业的分析建议"
        })
        
        if need_analysis.success:
            print(f"✅ 需求分析: {need_analysis.result['content'][:200]}...")
        
        # 步骤2: AI代码审查
        print("\n🔍 步骤2: AI代码审查")
        problem_code = """
def authenticate_user(username, password):
    # 直接SQL查询 - 存在注入风险
    query = f"SELECT * FROM users WHERE name='{username}' AND pwd='{password}'"
    user = database.execute(query).fetchone()
    
    if user:
        # 明文存储会话信息
        session_store[username] = {"logged_in": True, "role": user[2]}
        return True
    return False
"""
        
        code_review = await self.client.call_tool("ai_code_review", {
            "code": problem_code,
            "language": "python",
            "focus": "安全性和性能"
        })
        
        if code_review.success:
            print(f"✅ 代码审查: {code_review.result['content'][:200]}...")
        
        # 步骤3: AI生成改进文档
        print("\n📚 步骤3: AI生成改进文档")
        improved_code = """
import hashlib
import secrets
from sqlalchemy import text

def authenticate_user(username: str, password: str) -> Optional[Dict]:
    # 使用参数化查询防止SQL注入
    query = text("SELECT id, username, password_hash, role FROM users WHERE username = :username")
    user = db.execute(query, {"username": username}).fetchone()
    
    if user and verify_password(password, user.password_hash):
        # 生成安全的会话token
        session_token = secrets.token_urlsafe(32)
        # 存储会话信息到安全存储
        secure_session_store.set(session_token, {
            "user_id": user.id,
            "username": user.username,
            "role": user.role,
            "expires_at": datetime.now() + timedelta(hours=24)
        })
        return {"session_token": session_token}
    return None
"""
        
        documentation = await self.client.call_tool("ai_documentation", {
            "content": improved_code,
            "doc_type": "安全改进文档",
            "style": "详细"
        })
        
        if documentation.success:
            print(f"✅ 改进文档: {documentation.result['content'][:200]}...")
        
        # 步骤4: AI总结工作成果
        print("\n📄 步骤4: AI总结工作成果")
        summary_text = f"""
        用户认证系统优化项目完成：
        
        原始问题：{need_analysis.result['content'][:100] if need_analysis.success else '需求分析'}
        
        代码审查发现：{code_review.result['content'][:100] if code_review.success else '代码问题'}
        
        改进方案：{documentation.result['content'][:100] if documentation.success else '改进建议'}
        """
        
        final_summary = await self.client.call_tool("ai_summarize", {
            "text": summary_text,
            "length": "中等",
            "focus": "技术改进和安全提升"
        })
        
        if final_summary.success:
            print(f"✅ 工作总结: {final_summary.result['content']}")
        
        print("\n🎉 AI工具综合工作流演示完成！")
        print("展示了AI工具在实际开发场景中的协同应用")


async def main():
    """主函数"""
    tester = AIToolsTester()
    await tester.run_ai_tests()


if __name__ == "__main__":
    asyncio.run(main()) 